package com.lq.gitee;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class PreProcessor {

    public static final int BLACK = Color.BLACK.getRGB();

    public static final int WHITE = Color.WHITE.getRGB();

    private static void cropAndBinarization() throws IOException
    {
        for (File imageFile : Configuration.BILL_SRC.listFiles()) {
            BufferedImage src = ImageIO.read(imageFile);
            // 裁剪
            int cropWidth = 400;
            int cropHeight = src.getHeight() - 260;
            int startX = 340;
            int startY = 260;
            BufferedImage cropImage = new BufferedImage(cropWidth, cropHeight, BufferedImage.TYPE_INT_RGB);
            Graphics2D tg = cropImage.createGraphics();
            tg.drawImage(src,
                    0, 0, cropWidth, cropHeight,
                    startX, startY, startX + cropWidth, startY + cropHeight,
                    null);
            tg.dispose();
            // 二值化处理
            for (int y = 0; y < cropHeight; y++) {
                for (int x = 0; x < cropWidth; x++) {
                    int rgb = cropImage.getRGB(x, y);
                    int value = 0xff000000 | rgb;
                    int r = (value >> 16) & 0xFF;
                    int g = (value >> 8) & 0xFF;
                    int b = (value >> 0) & 0xFF;
                    boolean isBlack = r <= 100 && g <= 100 && b <= 100;
                    cropImage.setRGB(x, y, isBlack ? BLACK : WHITE);
                }
            }
            // 写入到预处理后的文件夹
            File cropFile = new File(Configuration.CROP_FOLDER, imageFile.getName());
            ImageIO.write(cropImage, "png", cropFile);
        }
    }

    private static class Partition {
        int s;
        int e;
        int length;
    }

    private static List<Partition> partition(int[] blackCounts)
    {
        List<Partition> result = new ArrayList<>();
        for (int i = 0, len = blackCounts.length; i < len; i++) {
            if (blackCounts[i] != 0) {
                for (int j = i + 1; j < len; j++) {
                    if (blackCounts[j] == 0) {
                        Partition p = new Partition();
                        p.s = i;
                        p.e = j - 1;
                        p.length = p.e - p.s + 1;
                        result.add(p);
                        i = j;
                        break;
                    }
                }
            }
        }
        return result;
    }

    private static void cropMoneyValue() throws IOException
    {
        int counter = 0;
        for (File imageFile : Configuration.CROP_FOLDER.listFiles()) {
            BufferedImage image = ImageIO.read(imageFile);
            int h = image.getHeight();
            int w = image.getWidth();
            int[] blackCounts = new int[h];
            for (int y = 0; y < h; y++) {
                int blackCount = 0;
                for (int x = 0; x < w; x++) {
                    int rgb = image.getRGB(x, y);
                    if (rgb == BLACK) {
                        ++blackCount;
                    }
                }
                blackCounts[y] = blackCount;
            }
            List<Partition> ps = partition(blackCounts);
            for (Partition p : ps) {
                if (p.length > 90) {
                    BufferedImage itemImage = new BufferedImage(w, p.length, BufferedImage.TYPE_INT_RGB);
                    Graphics2D g = itemImage.createGraphics();
                    g.drawImage(image,
                            0, 0, w, p.length,
                            0, p.s, w, p.e,
                            null);
                    File file = new File(Configuration.MONEY_VALUE_FOLDER, counter++ + ".png");
                    ImageIO.write(itemImage, "png", file);
                }
            }
        }
    }

    public static List<BufferedImage> toNumbers(File imageFile) throws IOException
    {
        List<BufferedImage> result = new ArrayList<>();

        BufferedImage image = ImageIO.read(imageFile);
        int h = image.getHeight();
        int w = image.getWidth();
        int[] blackCounts = new int[w];
        for (int x = 0; x < w; x++) {
            int blackCount = 0;
            for (int y = 0; y < h; y++) {
                int rgb = image.getRGB(x, y);
                if (rgb == BLACK) {
                    ++blackCount;
                }
            }
            blackCounts[x] = blackCount;
        }
        List<Partition> ps = partition(blackCounts);
        for (int i = 1; i < ps.size(); i++) {
            Partition p = ps.get(i);
            BufferedImage itemImage = new BufferedImage(p.length, h, BufferedImage.TYPE_INT_RGB);
            Graphics2D g = itemImage.createGraphics();
            g.drawImage(image,
                    0, 0, p.length, h,
                    p.s, 0, p.e, h,
                    null);
            result.add(itemImage);
        }
        return result;
    }

    private static void printNumbers() throws IOException
    {
        int counter = 0;
        for (File imageFile : Configuration.MONEY_VALUE_FOLDER.listFiles()) {
            List<BufferedImage> images = toNumbers(imageFile);
            for (BufferedImage image : images) {
                File file = new File(Configuration.NUMBERS_FOLDER, counter++ + ".png");
                ImageIO.write(image, "png", file);
            }
        }
    }

    public static void main(String[] args) throws IOException
    {
        cropAndBinarization();
        cropMoneyValue();
        printNumbers();
    }

}
